import os
import time
import logging
from typing import Dict, List, Optional


class WebBrowserTool:
    def __init__(self, headless=True, max_retries=3, lazy_init=True):
        self.headless = headless
        self.max_retries = max_retries
        self.lazy_init = lazy_init
        self.driver = None
        self.initialized = False
        self.logger = logging.getLogger(__name__)
        
        # 添加停止标志支持
        self._stop_flags = {}  # 存储每个会话的停止标志
        
        # 如果不是延迟初始化，立即初始化
        if not lazy_init:
            self._initialize_driver()

    def _find_chrome_binary(self):
        """查找Chrome浏览器可执行文件路径"""
        # Docker环境中的Chrome路径（根据Dockerfile配置）
        docker_paths = [
            "/usr/local/bin/google-chrome",
            "/usr/local/bin/chrome", 
            "/opt/chrome-linux64/chrome",
            "/opt/chrome/chrome"
        ]
        
        # Windows环境中的Chrome路径
        windows_paths = [
            "C:\\Program Files\\Google\\Chrome\\Application\\chrome.exe",
            "C:\\Program Files (x86)\\Google\\Chrome\\Application\\chrome.exe",
            os.path.expanduser("~\\AppData\\Local\\Google\\Chrome\\Application\\chrome.exe")
        ]
        
        # Linux/Mac环境中的Chrome路径
        linux_paths = [
            "/usr/bin/google-chrome",
            "/usr/bin/google-chrome-stable", 
            "/usr/bin/chromium-browser",
            "/usr/bin/chromium",
            "/snap/bin/chromium"
        ]
        
        # 根据操作系统选择搜索路径
        if os.name == 'nt':  # Windows
            search_paths = windows_paths + docker_paths
        else:  # Linux/Mac/Docker
            search_paths = docker_paths + linux_paths
            
        # 检查每个可能的路径
        for path in search_paths:
            if os.path.isfile(path) and os.access(path, os.X_OK):
                return path
                
        # 尝试通过which命令查找
        try:
            import subprocess
            result = subprocess.run(['which', 'google-chrome'], 
                                  capture_output=True, text=True, timeout=5)
            if result.returncode == 0 and result.stdout.strip():
                return result.stdout.strip()
                
            result = subprocess.run(['which', 'chrome'], 
                                  capture_output=True, text=True, timeout=5)
            if result.returncode == 0 and result.stdout.strip():
                return result.stdout.strip()
        except:
            pass
            
        return None

    def _find_chromedriver(self):
        """查找ChromeDriver可执行文件路径"""
        # Docker环境中的ChromeDriver路径（根据Dockerfile配置）
        docker_paths = [
            "/usr/local/bin/chromedriver"
        ]
        
        # Windows环境中的ChromeDriver路径
        windows_paths = [
            "chromedriver.exe",
            "C:\\chromedriver\\chromedriver.exe"
        ]
        
        # Linux/Mac环境中的ChromeDriver路径
        linux_paths = [
            "/usr/bin/chromedriver",
            "/usr/local/bin/chromedriver"
        ]
        
        # 根据操作系统选择搜索路径
        if os.name == 'nt':  # Windows
            search_paths = windows_paths + docker_paths
        else:  # Linux/Mac/Docker
            search_paths = docker_paths + linux_paths
            
        # 检查每个可能的路径
        for path in search_paths:
            if os.path.isfile(path) and os.access(path, os.X_OK):
                return path
                
        # 尝试通过which命令查找
        try:
            import subprocess
            result = subprocess.run(['which', 'chromedriver'], 
                                  capture_output=True, text=True, timeout=5)
            if result.returncode == 0 and result.stdout.strip():
                return result.stdout.strip()
        except:
            pass
            
        return None

    def _is_docker_environment(self):
        """检测是否在Docker环境中运行"""
        # 检查/.dockerenv文件
        if os.path.exists('/.dockerenv'):
            return True
            
        # 检查/proc/1/cgroup文件
        try:
            with open('/proc/1/cgroup', 'r') as f:
                content = f.read()
                if 'docker' in content or 'containerd' in content:
                    return True
        except:
            pass
            
        # 检查环境变量
        if os.environ.get('DOCKER_CONTAINER') or os.environ.get('container'):
            return True
            
        return False

    def _initialize_driver(self):
        """初始化浏览器驱动"""
        if self.initialized:
            return True

        self.logger.info("🚀 开始初始化浏览器驱动...")
        start_time = time.time()

        # 尝试初始化，带重试机制
        for attempt in range(self.max_retries):
            try:
                attempt_start = time.time()
                self._setup_driver()
                self.initialized = True
                elapsed = time.time() - start_time
                attempt_elapsed = time.time() - attempt_start
                self.logger.info(f"✅ 网页浏览器初始化成功 (第{attempt + 1}次尝试, 本次耗时: {attempt_elapsed:.2f}s, 总耗时: {elapsed:.2f}s)")
                return True
            except Exception as e:
                attempt_elapsed = time.time() - attempt_start
                self.logger.warning(f"⚠️ 网页浏览器第{attempt + 1}次初始化失败 (耗时: {attempt_elapsed:.2f}s): {e}")
                
                # 清理失败的驱动实例
                if hasattr(self, 'driver') and self.driver:
                    try:
                        self.driver.quit()
                    except:
                        pass
                    self.driver = None
                
                if attempt < self.max_retries - 1:
                    wait_time = min(2 ** attempt, 10)  # 指数退避，最大10秒
                    self.logger.info(f"⏳ 等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                else:
                    total_elapsed = time.time() - start_time
                    self.logger.error(f"❌ 网页浏览器初始化最终失败 (总耗时: {total_elapsed:.2f}s)，切换到模拟模式")
                    print(f"❌ 网页浏览器初始化最终失败 (总耗时: {total_elapsed:.2f}s)，切换到模拟模式")
                    self._setup_mock_driver()
                    return False

    def _ensure_initialized(self):
        """确保浏览器已初始化"""
        if not self.initialized and self.lazy_init:
            self.logger.info("🚀 首次使用浏览器功能，开始初始化...")
            print("🚀 首次使用浏览器功能，开始初始化...")
            return self._initialize_driver()
        return self.initialized

    def _setup_driver(self):
        """设置真实的浏览器驱动"""
        try:
            from selenium import webdriver
            from selenium.webdriver.chrome.options import Options
            from selenium.webdriver.chrome.service import Service
            from selenium.webdriver.common.by import By
            from selenium.webdriver.support.ui import WebDriverWait
            from selenium.webdriver.support import expected_conditions as EC
            from selenium.common.exceptions import TimeoutException, WebDriverException

            chrome_options = Options()

            if self.headless:
                chrome_options.add_argument("--headless=new")  # 使用新的headless模式

            # 基础安全参数
            chrome_options.add_argument("--no-sandbox")
            chrome_options.add_argument("--disable-dev-shm-usage")
            chrome_options.add_argument("--disable-gpu")
            chrome_options.add_argument("--window-size=1280,720")

            # Docker环境特殊参数
            chrome_options.add_argument("--disable-setuid-sandbox")
            chrome_options.add_argument("--disable-background-timer-throttling")
            chrome_options.add_argument("--disable-backgrounding-occluded-windows")
            chrome_options.add_argument("--disable-renderer-backgrounding")
            chrome_options.add_argument("--disable-features=TranslateUI")
            chrome_options.add_argument("--disable-ipc-flooding-protection")

            # 加速启动的参数
            chrome_options.add_argument("--disable-extensions")
            chrome_options.add_argument("--disable-plugins")
            chrome_options.add_argument("--disable-images")  # 禁用图片加载
            chrome_options.add_argument("--disable-javascript")  # 对于简单搜索可以禁用JS
            chrome_options.add_argument("--disable-css")  # 禁用CSS加载
            chrome_options.add_argument("--disable-web-security")
            chrome_options.add_argument("--disable-client-side-phishing-detection")
            chrome_options.add_argument("--disable-sync")
            chrome_options.add_argument("--disable-default-apps")
            chrome_options.add_argument("--no-first-run")
            chrome_options.add_argument("--no-default-browser-check")
            chrome_options.add_argument("--disable-logging")
            chrome_options.add_argument("--disable-notifications")
            chrome_options.add_argument("--disable-permissions-api")

            # 内存和性能优化
            chrome_options.add_argument("--memory-pressure-off")
            chrome_options.add_argument("--max_old_space_size=4096")
            # 移除单进程模式，添加稳定性参数
            chrome_options.add_argument("--disable-crash-reporter")
            chrome_options.add_argument("--disable-in-process-stack-traces")
            chrome_options.add_argument("--disable-logging")
            chrome_options.add_argument("--disable-dev-tools")
            chrome_options.add_argument("--log-level=3")  # 只显示致命错误

            # 检测环境并设置Chrome路径
            chrome_binary_path = self._find_chrome_binary()
            if chrome_binary_path:
                chrome_options.binary_location = chrome_binary_path
                self.logger.info(f"✅ 找到Chrome浏览器: {chrome_binary_path}")
            else:
                raise Exception("Chrome浏览器未安装")

            # 设置ChromeDriver路径
            chromedriver_path = self._find_chromedriver()
            if chromedriver_path:
                service = Service(executable_path=chromedriver_path)
                self.logger.info(f"✅ 找到ChromeDriver: {chromedriver_path}")
            else:
                # 尝试使用webdriver-manager作为备选
                try:
                    from webdriver_manager.chrome import ChromeDriverManager
                    service = Service(ChromeDriverManager().install())
                    self.logger.info("✅ 使用webdriver-manager自动下载ChromeDriver")
                except Exception as e:
                    self.logger.warning(f"⚠️ webdriver-manager失败: {e}")
                    service = Service()  # 使用默认路径

            # Windows环境特殊设置
            if os.name == 'nt':
                service.creation_flags = 0x08000000  # CREATE_NO_WINDOW flag for Windows

            # 添加启动超时控制
            try:
                self.driver = webdriver.Chrome(service=service, options=chrome_options)
            except Exception as e:
                if "timeout" in str(e).lower():
                    raise TimeoutException(f"Chrome启动超时: {e}")
                raise
            
            # 设置页面加载超时
            self.driver.set_page_load_timeout(30)  # 增加到30秒
            self.driver.implicitly_wait(10)  # 增加到10秒

            # 简单测试连接，添加超时控制
            try:
                self.driver.get("about:blank")
            except TimeoutException:
                self.logger.warning("测试页面加载超时，但驱动可能仍可用")

            self.logger.info("✅ Chrome浏览器驱动设置成功")

        except TimeoutException as e:
            self.logger.error(f"❌ 浏览器启动超时: {e}")
            raise
        except WebDriverException as e:
            self.logger.error(f"❌ WebDriver异常: {e}")
            raise
        except ImportError as e:
            self.logger.error(f"❌ 缺少必要的依赖: {e}")
            raise
        except Exception as e:
            self.logger.error(f"❌ 浏览器驱动设置失败: {e}")
            raise

    def _setup_mock_driver(self):
        """设置模拟浏览器驱动"""
        self.driver = None
        self.initialized = False
        self.logger.info("🔄 切换到模拟浏览器模式")

    def search_web(self, query: str, max_results: int = 5, session_id: str = None) -> Dict:
        """网页搜索"""
        # 确保浏览器已初始化
        if not self._ensure_initialized():
            return {
                "error": "网页浏览器初始化失败",
                "results": [],
                "note": "浏览器功能暂不可用，已切换到模拟模式"
            }

        if not self.driver:
            return self._get_mock_search_results(query, max_results)

        try:
            from selenium.webdriver.common.by import By
            from selenium.webdriver.support.ui import WebDriverWait
            from selenium.webdriver.support import expected_conditions as EC
            from urllib.parse import quote

            # 检查停止标志
            if session_id and self._stop_flags.get(session_id, False):
                return {
                    "error": "搜索已被停止",
                    "results": [],
                    "stopped": True
                }

            # 使用百度搜索，添加重试机制
            search_url = f"https://www.baidu.com/s?wd={quote(query)}"
            
            # 重试机制
            max_retries = 3
            for attempt in range(max_retries):
                try:
                    self.driver.get(search_url)
                    break
                except TimeoutException as e:
                    if attempt == max_retries - 1:
                        raise TimeoutException(f"页面加载超时，已重试{max_retries}次: {e}")
                    self.logger.warning(f"页面加载超时，正在重试 ({attempt + 1}/{max_retries})")
                    time.sleep(2)  # 等待2秒后重试

            # 等待结果加载，增加超时时间
            WebDriverWait(self.driver, 20).until(
                EC.presence_of_element_located((By.TAG_NAME, "body"))
            )

            # 提取搜索结果
            results = []
            try:
                search_elements = self.driver.find_elements(By.CSS_SELECTOR, ".result")
                for i, element in enumerate(search_elements[:max_results]):
                    # 在每次循环中检查停止标志
                    if session_id and self._stop_flags.get(session_id, False):
                        return {
                            "error": "搜索已被停止",
                            "results": results,  # 返回已获取的部分结果
                            "stopped": True
                        }
                    
                    try:
                        title_element = element.find_element(By.CSS_SELECTOR, "h3 a")
                        title = title_element.text
                        url = title_element.get_attribute("href")

                        # 尝试获取摘要
                        try:
                            snippet_element = element.find_element(By.CSS_SELECTOR, ".c-abstract")
                            snippet = snippet_element.text
                        except:
                            snippet = ""

                        results.append({
                            "title": title,
                            "url": url,
                            "snippet": snippet,
                            "rank": i + 1
                        })
                    except Exception as e:
                        continue
            except Exception as e:
                self.logger.warning(f"提取搜索结果失败: {e}")

            return {
                "query": query,
                "results": results,
                "search_engine": "百度",
                "result_count": len(results)
            }

        except TimeoutException as e:
            self.logger.error(f"搜索页面加载超时: {e}")
            return {
                "error": f"搜索超时: 页面加载时间过长，请稍后重试",
                "results": [],
                "timeout": True
            }
        except Exception as e:
            self.logger.error(f"搜索过程中出错: {e}")
            return {
                "error": f"搜索失败: {str(e)}",
                "results": []
            }

    def get_page_content(self, url: str, timeout: int = 10, session_id: str = None) -> Dict:
        """获取网页内容"""
        # 检查停止标志
        if session_id and self._stop_flags.get(session_id, False):
            return {
                "error": "用户停止了搜索",
                "content": "",
                "note": "搜索被用户停止"
            }
            
        # 确保浏览器已初始化
        if not self._ensure_initialized():
            return {
                "error": "网页浏览器初始化失败",
                "content": "",
                "note": "浏览器功能暂不可用"
            }

        if not self.driver:
            return {
                "error": "浏览器驱动不可用",
                "content": "",
                "note": "无法获取页面内容"
            }

        try:
            from selenium.webdriver.common.by import By
            from selenium.webdriver.support.ui import WebDriverWait
            from selenium.webdriver.support import expected_conditions as EC

            # 在导航前检查停止标志
            if session_id and self._stop_flags.get(session_id, False):
                return {
                    "error": "用户停止了搜索",
                    "content": "",
                    "note": "搜索被用户停止"
                }

            self.driver.get(url)
            
            # 在等待页面加载时检查停止标志
            wait = WebDriverWait(self.driver, timeout)
            wait.until(EC.presence_of_element_located((By.TAG_NAME, "body")))
            
            # 在获取内容前检查停止标志
            if session_id and self._stop_flags.get(session_id, False):
                return {
                    "error": "用户停止了搜索",
                    "content": "",
                    "note": "搜索被用户停止"
                }

            title = self.driver.title
            body = self.driver.find_element(By.TAG_NAME, "body")
            content = body.text[:2000]  # 限制内容长度

            return {
                "url": url,
                "title": title,
                "content": content + "..." if len(content) == 2000 else content,
                "success": True
            }

        except Exception as e:
            self.logger.error(f"获取页面内容失败: {e}")
            return {
                "error": f"获取页面内容失败: {str(e)}",
                "url": url
            }

    def search_web_with_fallback(self, query: str, max_results: int = 5) -> Dict:
        """带降级方案的网页搜索"""
        # 先尝试真实搜索
        result = self.search_web(query, max_results)

        # 如果失败或没有结果，返回模拟结果
        if "error" in result or not result.get("results"):
            return self._get_mock_search_results(query, max_results)

        return result

    def _get_mock_search_results(self, query: str, max_results: int) -> Dict:
        """获取模拟搜索结果"""
        mock_results = [
            {
                "title": f"关于 {query} 的搜索结果 1",
                "url": "https://example.com/result1",
                "snippet": f"这是关于 {query} 的模拟搜索结果。实际搜索功能由于浏览器配置问题暂时不可用。",
                "rank": 1
            },
            {
                "title": f"关于 {query} 的搜索结果 2",
                "url": "https://example.com/result2",
                "snippet": f"更多关于 {query} 的信息。请检查浏览器驱动配置。",
                "rank": 2
            }
        ]

        return {
            "query": query,
            "results": mock_results[:max_results],
            "search_engine": "模拟模式",
            "result_count": len(mock_results[:max_results]),
            "note": "⚠️ 当前为模拟模式，实际搜索功能需要浏览器支持"
        }

    def close(self):
        """关闭浏览器"""
        if self.driver:
            try:
                self.driver.quit()
            except:
                pass
            self.driver = None

    def stop_generation(self, session_id: str) -> bool:
        """停止指定会话的生成"""
        try:
            self._stop_flags[session_id] = True
            return True
        except Exception as e:
            self.logger.error(f"停止生成失败: {e}")
            return False

    def stop_all_generations(self) -> bool:
        """停止所有活跃的生成"""
        try:
            for session_id in list(self._stop_flags.keys()):
                self._stop_flags[session_id] = True
            return True
        except Exception as e:
            self.logger.error(f"停止所有生成失败: {e}")
            return False

    def clear_stop_flag(self, session_id: str):
        """清理指定会话的停止标志"""
        if session_id in self._stop_flags:
            del self._stop_flags[session_id]

    def __del__(self):
        """析构函数确保浏览器关闭"""
        self.close()
